#!/bin/sh

ACTION=$1
OVERRIDE_ENCODER=$2

export DISPLAY=$(getLocalXDisplay)

RESOLUTION=$(batocera-resolution currentResolution)
REFRESH_RATE=$(batocera-resolution refreshRate)
OUTPUT="/userdata/recordings/capture-$(date +%Y.%m.%d-%Hh%M.%S).mkv"
VAAPI_DEVICE="/dev/dri/renderD128"  # Default VAAPI device ?
PULSEAUDIO_MONITOR=$(LC_ALL=C pactl info | grep -E '^Default Sink: ' | sed -e s+"^Default Sink: "+""+ | head -1).monitor

mkdir -p /userdata/recordings || exit 1

# Check for a NVIDIA GPU in the system
has_nvidia_gpu() {
    NVIDIA_DEV=$(lspci -mn | awk '{ gsub("\"",""); if (($2 ~ "030[0-2]") && ($3 == "10de" || $3 == "12d2")) { print $1 } }')
    
    if [ -n "$NVIDIA_DEV" ]; then
        return 0
    else
        return 1
    fi
}

# Check for hardware encoding support
check_hardware_encoding() {
    if [ -n "$OVERRIDE_ENCODER" ]; then
        case "$OVERRIDE_ENCODER" in
            hevc)
                if has_nvidia_gpu; then
                    ENCODER="hevc_nvenc"
                    echo "Using user-selected HEVC hardware encoding via NVENC."
                else
                    ENCODER="hevc_vaapi"
                    echo "Using user-selected HEVC hardware encoding via VA-API."
                fi
                ;;
            h264)
                if has_nvidia_gpu; then
                    ENCODER="h264_nvenc"
                    echo "Using user-selected H.264 hardware encoding via NVENC."
                else
                    ENCODER="h264_vaapi"
                    echo "Using user-selected H.264 hardware encoding via VA-API."
                fi
                ;;
            mpeg2)
                if has_nvidia_gpu; then
                    echo "Nvidia GPU's don't support hardware MPEG-2 encoding."
                    echo "Falling back to software H.264 encoding."
                    ENCODER="libx264"
                else
                    ENCODER="mpeg2_vaapi"
                    echo "Using user-selected MPEG-2 hardware encoding via VA-API."
                fi
                ;;
            software)
                ENCODER="libx264"
                echo "Using user-selected software H.264 encoding."
                ;;
            *)
                echo "Invalid encoder option. Using auto-detection."
                auto_detect_encoder
                ;;
        esac
    else
        auto_detect_encoder
    fi

    echo "With an audio input source of: ${PULSEAUDIO_MONITOR}"
    if [ "$ENCODER" = "libx264" ]; then
        echo "A resolution of: ${RESOLUTION}"
    else
        echo "A resolution of: ${RESOLUTION} and a framerate of: ${REFRESH_RATE}"       
    fi
    echo
    echo "Recording to file: ${OUTPUT}"
}

# Auto detect the encoder type, prefer VA-API
auto_detect_encoder() {
    if vainfo 2>/dev/null | grep -q "VAProfileHEVCMain.*VAEntrypointEncSlice"; then
        ENCODER="hevc_vaapi"
        echo
        echo "Using: HEVC hardware encoding via VA-API."
    elif has_nvidia_gpu; then
        ENCODER="hevc_nvenc"
        echo
        echo "Using: HEVC hardware encoding via NVENC."
    elif vainfo 2>/dev/null | grep -q "VAProfileH264Main.*VAEntrypointEncSlice"; then
        ENCODER="h264_vaapi"
        echo
        echo "Using: H.264 hardware encoding via VA-API."
    elif has_nvidia_gpu; then
        ENCODER="h264_nvenc"
        echo
        echo "Using: H.264 hardware encoding via NVENC."    
    elif vainfo 2>/dev/null | grep -q "VAProfileMPEG2Main.*VAEntrypointEncSlice"; then
        ENCODER="mpeg2_vaapi"
        echo
        echo "Using: MPEG-2 hardware encoding via VA-API."
    else
        ENCODER="libx264"  # Fallback to software H.264 encoding
        echo
        echo "Hardware encoding is not available."
        echo "Using: software encoder ${ENCODER}."
    fi
}

usage() {
    echo "Usage: ${1} [OPTION] [ENCODER]"
    echo "Options:"
    echo "  --default"
    echo "  --fast"
    echo "  --compress"
    echo "  --high-quality"
    echo "  --epic-mode"
    echo
    echo "Encoder options:"
    echo "  hevc      - HEVC hardware encoding"
    echo "  h264      - H.264 hardware encoding"
    echo "  mpeg2     - MPEG-2 hardware encoding (Non-Nvidia cards)"
    echo "  software  - Software H.264 encoding"
    echo
    echo "Example: ${1} --default hevc"
}

exit_info() {
    echo >&2
    echo "+-------------------------------+"   >&2
    echo "| Press Ctrl+c to end recording |"   >&2
    echo "+-------------------------------+"   >&2
    echo >&2
}

default_record(){
    check_hardware_encoding
    exit_info
    if [ "$ENCODER" = "hevc_vaapi" ] || [ "$ENCODER" = "h264_vaapi" ] || [ "$ENCODER" = "mpeg2_vaapi" ]; then
        ffmpeg -probesize 20M -video_size "${RESOLUTION}" -framerate "${REFRESH_RATE}" \
               -thread_queue_size 1024 -f x11grab -i :0.0+0,0 \
               -f pulse -ac 2 -channel_layout stereo -i "${PULSEAUDIO_MONITOR}" \
               -vaapi_device "${VAAPI_DEVICE}" \
               -vf 'format=nv12,hwupload' -c:v ${ENCODER} -b:v 5000k -c:a aac -async 1 "${OUTPUT}" \
               -loglevel warning
    elif [ "$ENCODER" = "hevc_nvenc" ] || [ "$ENCODER" = "h264_nvenc" ] || [ "$ENCODER" = "mpeg2_nvenc" ]; then
        ffmpeg -probesize 20M -video_size "${RESOLUTION}" -framerate "${REFRESH_RATE}" \
               -thread_queue_size 1024 -f x11grab -i :0.0+0,0 \
               -f pulse -ac 2 -channel_layout stereo -i "${PULSEAUDIO_MONITOR}" \
               -vf 'format=nv12' -c:v ${ENCODER} -b:v 5000k -c:a aac -async 1 "${OUTPUT}" \
               -loglevel warning
    else
        ffmpeg -probesize 20M -video_size "${RESOLUTION}" -framerate 25 \
               -thread_queue_size 1024 -f x11grab -i :0.0+0,0 \
               -f pulse -ac 2 -channel_layout stereo -i "${PULSEAUDIO_MONITOR}" \
               -c:v ${ENCODER} -b:v 5000k -c:a aac -async 1 "${OUTPUT}" \
               -loglevel warning
    fi
}

fast_record(){
    check_hardware_encoding
    exit_info
    if [ "$ENCODER" = "hevc_vaapi" ] || [ "$ENCODER" = "h264_vaapi" ] || [ "$ENCODER" = "mpeg2_vaapi" ]; then
        ffmpeg -probesize 20M -video_size "${RESOLUTION}" -framerate "${REFRESH_RATE}" \
               -thread_queue_size 1024 -f x11grab -i :0.0 \
               -f pulse -ac 2 -channel_layout stereo -i "${PULSEAUDIO_MONITOR}" \
               -vaapi_device "${VAAPI_DEVICE}" \
               -vf 'format=nv12,hwupload' -c:v ${ENCODER} -qp 0 -preset ultrafast -c:a aac -async 1 "${OUTPUT}" \
               -loglevel warning
    elif [ "$ENCODER" = "hevc_nvenc" ] || [ "$ENCODER" = "h264_nvenc" ] || [ "$ENCODER" = "mpeg2_nvenc" ]; then
        ffmpeg -probesize 20M -video_size "${RESOLUTION}" -framerate "${REFRESH_RATE}" \
               -thread_queue_size 1024 -f x11grab -i :0.0 \
               -f pulse -ac 2 -channel_layout stereo -i "${PULSEAUDIO_MONITOR}" \
               -vf 'format=nv12' -c:v ${ENCODER} -qp 0 -preset ultrafast -c:a aac -async 1 "${OUTPUT}" \
               -loglevel warning
    else
        ffmpeg -probesize 20M -video_size "${RESOLUTION}" -framerate 30 \
               -thread_queue_size 1024 -f x11grab -i :0.0 \
               -f pulse -ac 2 -channel_layout stereo -i "${PULSEAUDIO_MONITOR}" \
               -c:v ${ENCODER} -crf 0 -preset ultrafast -c:a aac -async 1 "${OUTPUT}" \
               -loglevel warning
    fi
}

compress_record(){
    check_hardware_encoding
    exit_info
    if [ "$ENCODER" = "hevc_vaapi" ] || [ "$ENCODER" = "h264_vaapi" ] || [ "$ENCODER" = "mpeg2_vaapi" ]; then
        ffmpeg -probesize 20M -video_size "${RESOLUTION}" -framerate "${REFRESH_RATE}" \
               -thread_queue_size 1024 -f x11grab -i :0.0 \
               -f pulse -ac 2 -channel_layout stereo -i "${PULSEAUDIO_MONITOR}" \
               -vaapi_device "${VAAPI_DEVICE}" \
               -vf 'format=nv12,hwupload' -c:v ${ENCODER} -crf 20 -preset veryfast -c:a aac -async 1 "${OUTPUT}" \
               -loglevel warning
    elif [ "$ENCODER" = "hevc_nvenc" ] || [ "$ENCODER" = "h264_nvenc" ] || [ "$ENCODER" = "mpeg2_nvenc" ]; then
        ffmpeg -probesize 20M -video_size "${RESOLUTION}" -framerate "${REFRESH_RATE}" \
               -thread_queue_size 1024 -f x11grab -i :0.0 \
               -f pulse -ac 2 -channel_layout stereo -i "${PULSEAUDIO_MONITOR}" \
               -vf 'format=nv12' -c:v ${ENCODER} -crf 20 -preset veryfast -c:a aac -async 1 "${OUTPUT}" \
               -loglevel warning
    else
        ffmpeg -probesize 20M -video_size "${RESOLUTION}" -framerate 30 \
               -thread_queue_size 1024 -f x11grab -i :0.0 \
               -f pulse -ac 2 -channel_layout stereo -i "${PULSEAUDIO_MONITOR}" \
               -c:v ${ENCODER} -crf 20 -preset veryfast -c:a aac -async 1 "${OUTPUT}" \
               -loglevel warning
    fi
}

high_quality_record(){
    check_hardware_encoding
    exit_info
    if [ "$ENCODER" = "hevc_vaapi" ] || [ "$ENCODER" = "h264_vaapi" ] || [ "$ENCODER" = "mpeg2_vaapi" ]; then
        ffmpeg -probesize 20M -video_size "${RESOLUTION}" -framerate "${REFRESH_RATE}" \
               -thread_queue_size 1024 -f x11grab -i :0.0 \
               -f pulse -ac 2 -channel_layout stereo -i "${PULSEAUDIO_MONITOR}" \
               -vaapi_device "${VAAPI_DEVICE}" \
               -vf 'format=nv12,hwupload' -c:v ${ENCODER} -qp 23 -preset fast -c:a aac -async 1 "${OUTPUT}" \
               -loglevel warning
    elif [ "$ENCODER" = "hevc_nvenc" ] || [ "$ENCODER" = "h264_nvenc" ] || [ "$ENCODER" = "mpeg2_nvenc" ]; then
        ffmpeg -probesize 20M -video_size "${RESOLUTION}" -framerate "${REFRESH_RATE}" \
               -thread_queue_size 1024 -f x11grab -i :0.0 \
               -f pulse -ac 2 -channel_layout stereo -i "${PULSEAUDIO_MONITOR}" \
               -vf 'format=nv12' -c:v ${ENCODER} -qp 23 -preset fast -c:a aac -async 1 "${OUTPUT}" \
               -loglevel warning
    else
        ffmpeg -probesize 20M -video_size "${RESOLUTION}" -framerate 60 \
               -thread_queue_size 1024 -f x11grab -i :0.0 \
               -f pulse -ac 2 -channel_layout stereo -i "${PULSEAUDIO_MONITOR}" \
               -c:v ${ENCODER} -qp 23 -preset fast -c:a aac -async 1 "${OUTPUT}" \
               -loglevel warning
    fi
}

epic_mode_record(){
    check_hardware_encoding
    exit_info
    if [ "$ENCODER" = "hevc_vaapi" ] || [ "$ENCODER" = "h264_vaapi" ] || [ "$ENCODER" = "mpeg2_vaapi" ]; then
        ffmpeg -probesize 20M -video_size "${RESOLUTION}" -framerate "${REFRESH_RATE}" \
               -thread_queue_size 1024 -f x11grab -i :0.0 \
               -f pulse -ac 2 -channel_layout stereo -i "${PULSEAUDIO_MONITOR}" \
               -vaapi_device "${VAAPI_DEVICE}" \
               -vf 'format=nv12,hwupload' -c:v ${ENCODER} -qp 18 -preset medium -c:a aac -async 1 "${OUTPUT}" \
               -loglevel warning
    elif [ "$ENCODER" = "hevc_nvenc" ] || [ "$ENCODER" = "h264_nvenc" ] || [ "$ENCODER" = "mpeg2_nvenc" ]; then
        ffmpeg -probesize 20M -video_size "${RESOLUTION}" -framerate "${REFRESH_RATE}" \
               -thread_queue_size 1024 -f x11grab -i :0.0 \
               -f pulse -ac 2 -channel_layout stereo -i "${PULSEAUDIO_MONITOR}" \
               -vf 'format=nv12' -c:v ${ENCODER} -qp 18 -preset medium -c:a aac -async 1 "${OUTPUT}" \
               -loglevel warning    
    else
        ffmpeg -probesize 20M -video_size "${RESOLUTION}" -framerate 60 \
               -thread_queue_size 1024 -f x11grab -i :0.0 \
               -f pulse -ac 2 -channel_layout stereo -i "${PULSEAUDIO_MONITOR}" \
               -c:v ${ENCODER} -qp 18 -preset medium -c:a aac -async 1 "${OUTPUT}" \
               -loglevel warning
    fi
}

case "${ACTION}" in
    --fast)
        fast_record || exit 1
        ;;
    --default)
        default_record || exit 1
        ;;
    --compress)
        compress_record || exit 1
        ;;
    --high-quality)
        high_quality_record || exit 1
        ;;
    --epic-mode)
        epic_mode_record || exit 1
        ;;
    -h|--help)
        usage "${0}" || exit 1
        exit 0
        ;;
    *)
        default_record || exit 1
        ;;
esac
